//
// Created by Sorel on 2021/3/29.
//

#ifndef BASIC_EAS_PSO_HPP
#define BASIC_EAS_PSO_HPP

namespace PSO {

    class Particle {
    public:
        vector<double> position;  // current position
        vector<double> velocity;  // current velocity
        double fitness;  // current fitness
        vector<double> pbest_position;  // particle's personal best position in history
        double pbest_fitness;  // particle's personal best fitness in history
    };

    class BasicPSO {
    private:
        int pop_size;
        int dimension;
        double lower_bound, upper_bound;
        double w, c1, c2;

        vector<Particle> population;
        vector<double> gbest_position;
        double gbest_fitness;
        void print_pop(bool show_pos) const;

        /** the basic PSO operators **/
        void evaluate(int func_num);
        void swarm_evolve();

        string save_path;
        vector<double> value_list;  // record the history of gbest fitness
        void save_results();
    public:
        void set_save_path(const string& _save_path);
        void init(int _pop_size, int _dimension, double _inertia, double _c1, double _c2);
        void set_bounds(double _lower, double _upper);
        void execute(int func_num, int depth, int trial);

    };

    void BasicPSO::init(int _pop_size, int _dimension, double _inertia, double _c1, double _c2) {
        this->pop_size = _pop_size;
        this->dimension = _dimension;
        this->w = _inertia;
        this->c1 = _c1;
        this->c2 = _c2;
        population.clear();
        population.resize(pop_size);
        for (auto &p : population) {
            p.position.resize(dimension);
            p.velocity.resize(dimension);
            p.pbest_position.resize(dimension);
        }
        gbest_position.resize(dimension);
    }

    void BasicPSO::set_bounds(double _lower, double _upper) {
        this->lower_bound = _lower;
        this->upper_bound = _upper;
    }

    void BasicPSO::execute(int func_num, int depth, int trial) {
        value_list.clear();
        value_list.resize(depth);
        fill(value_list.begin(), value_list.end(), 0);
        /** the main PSO procedure starts **/
        int _trial = 0;
        cout<<"trial ";
        int i;
        while (_trial < trial) {
            // initialization of population
            for (auto &p : population) {
                for (i = 0; i < dimension; ++i) {
                    p.pbest_position[i] = p.position[i] = random_double(lower_bound, upper_bound);
                    p.velocity[i] = 0.0;  // zero initial velocity
//                    p.velocity[i] = random_double(lower_bound, upper_bound);  // random initial velocity
                    p.pbest_fitness = numeric_limits<double>::max();  // regarding to minimization problem
                }
            }
            gbest_fitness = numeric_limits<double>::max();  // regarding to minimization problem
            evaluate(func_num);

            // swarm evolution
            int d = 0;
            while (d < depth) {
                w = 0.9 - 0.5 * (double(d) / double(depth));
                swarm_evolve();
                evaluate(func_num);
                value_list[d] += gbest_fitness;
                d++;
            }
            _trial++;
            cout<<_trial<<", ";
        }
        /** the main PSO procedure ends **/
        for (auto &value : value_list) {
            value = value / trial;
        }
        save_results();
    }

    void BasicPSO::evaluate(int func_num) {
        for (int p = 0; p < pop_size; ++p) {
            /** replace your evaluation operator here **/
            cec14_test_func(population[p].position.data(), &population[p].fitness, dimension, 1, func_num);

            if (population[p].pbest_fitness > population[p].fitness) {  // update pbest
                population[p].pbest_fitness = population[p].fitness;
                population[p].pbest_position.assign(population[p].position.begin(), population[p].position.end());
            }
            if (gbest_fitness > population[p].fitness) {  // update gbest
                gbest_fitness = population[p].fitness;
                gbest_position.assign(population[p].position.begin(), population[p].position.end());
            }
        }
    }

    void BasicPSO::swarm_evolve() {
        int i;
        for (auto & p : population) {
            // update velocity
            for (i = 0; i < dimension; ++i) {
                p.velocity[i] = w * p.velocity[i]
                                + c1 * random_p() * (p.pbest_position[i] - p.position[i])
                                + c2 * random_p() * (gbest_position[i] - p.position[i]);
            }
            // update position
            for (i = 0; i < dimension; ++i) {
                p.position[i] = p.position[i] + p.velocity[i];
                // correct position by reflection
                if (p.position[i] < lower_bound) {
                    p.position[i] = lower_bound + (lower_bound - p.position[i]);
                }
                if (p.position[i] > upper_bound) {
                    p.position[i] = upper_bound - (p.position[i] - upper_bound);
                }
            }
        }
    }

    void BasicPSO::set_save_path(const string &_save_path) {
        this->save_path = _save_path;
        string file_name = save_path;
        ofstream out_file(file_name);
        if (!out_file.is_open()) {
            cout<<"save_path invalid, please check directories. "<<endl;
            exit(-1);
        }
        out_file.close();
    }

    void BasicPSO::save_results() {
        string file_name = save_path;
        ofstream out_file(file_name);
        if (!out_file.is_open()) {
            cout<<"save_path invalid, please check directories. "<<endl;
            exit(-1);
        }
        for (double value : value_list) {
            out_file<<value<<endl;
        }
        out_file.close();
        cout<<save_path<<" created. "<<endl;
    }

    void BasicPSO::print_pop(bool show_pos) const {
        for (int p = 0; p < pop_size; ++p) {
            cout<<"p "<<p<<": ";
            if (show_pos) {
                for (int i = 0; i < dimension; ++i) {
                    cout<<population[p].position[i]<<", ";
                }
            } cout<<"fitness: "<<population[p].fitness<<". "<<endl;
        }
        cout<<"gbest: ";
        if (show_pos) {
            for (int i = 0; i < dimension; ++i) {
                cout<<gbest_position[i]<<", ";
            }
        } cout<<"fitness: "<<gbest_fitness<<". "<<endl;
    }

}

#endif //BASIC_EAS_PSO_HPP
